package cn.com.basic.framework.aliyun.odps.utils;

import cn.com.basic.framework.aliyun.odps.error.OdpsError;
import cn.com.basic.framework.exception.exceptions.SysException;
import com.aliyun.odps.TableSchema;
import com.aliyun.odps.data.Record;
import com.aliyun.odps.tunnel.TableTunnel;
import com.aliyun.odps.tunnel.io.ProtobufRecordPack;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;

public class OdpsWriterProxy {
    private static final Logger LOG = LoggerFactory
            .getLogger(OdpsWriterProxy.class);

    private volatile boolean printColumnLess;// 是否打印对于源头字段数小于 ODPS 目的表的行的日志

    private TableTunnel.UploadSession slaveUpload;
    private TableSchema schema;
    private int maxBufferSize;
    private ProtobufRecordPack protobufRecordPack;
    private int protobufCapacity;
    private AtomicLong blockId;

    private boolean isCompress;

    public OdpsWriterProxy(TableTunnel.UploadSession slaveUpload, int blockSizeInMB, AtomicLong blockId, boolean isCompress) throws IOException {
        this.slaveUpload = slaveUpload;
        this.schema = this.slaveUpload.getSchema();

        this.blockId = blockId;
        this.isCompress = isCompress;

        // 初始化与 buffer 区相关的值
        this.maxBufferSize = (blockSizeInMB - 4) * 1024 * 1024;
        this.protobufCapacity = blockSizeInMB * 1024 * 1024;
        this.protobufRecordPack = new ProtobufRecordPack(this.schema, null, this.protobufCapacity);
        printColumnLess = true;

    }

    public long writeOneRecord(Record record, List<Long> blocks) {
        try {
            if (null == record) {
                return 0;
            }
            protobufRecordPack.append(record);
            if (protobufRecordPack.getProtobufStream().size() >= maxBufferSize) {
                long startTimeInNs = System.nanoTime();
                OdpsUtil.slaveWriteOneBlock(this.slaveUpload,
                        protobufRecordPack, blockId.get(), this.isCompress);
                LOG.info("write block {} ok.", blockId.get());
                blocks.add(blockId.get());
                protobufRecordPack.reset();
                this.blockId.incrementAndGet();
                return System.nanoTime() - startTimeInNs;
            }
        } catch (Exception e) {
            throw SysException.newException(OdpsError.ODPS_WRITER_ERROR, e.getMessage());
        }
        return 0;
    }

    public long writeRemainingRecord(List<Long> blocks) {
        try {
            if (protobufRecordPack.getProtobufStream().size() != 0) {
                long startTimeInNs = System.nanoTime();
                OdpsUtil.slaveWriteOneBlock(this.slaveUpload,
                        protobufRecordPack, blockId.get(), this.isCompress);
                LOG.info("write remain block {} ok.", blockId.get());

                blocks.add(blockId.get());
                // reset the buffer for next block
                protobufRecordPack.reset();
                return System.nanoTime() - startTimeInNs;
            }
        } catch (IOException e) {
            throw SysException.newException(OdpsError.ODPS_WRITER_ERROR, e.getMessage());
        }
        return 0;
    }

}
